from model.ctc_model import model_all,conv_shape
from pretreat import img_pretreat
from PIL import Image
import numpy as np
from keras import backend as K
from config import Config
import threading
import os


def predict3(filename):
    '''
    根据识别结果返回颜色和字符返回3个字符的识别结果
    :param filename: 验证码图片路径
    :return:
    '''
    img = Image.open(filename)
    arr = np.expand_dims(img_pretreat(img),0)
    model_all.load_weights(Config.model_path3)
    result = model_all.predict(arr)
    pred_str = K.get_value(
        K.ctc_decode(result[0], input_length=np.ones(1, dtype=int) * int(conv_shape[1]), )[0][0])[:, :6]  # ctc解码，[:,:6]只取前6位
    pred_color = K.get_value(
        K.ctc_decode(result[1], input_length=np.ones(1, dtype=int) * int(conv_shape[1]), )[0][0])[:, :6]
    with open(Config.str_table_path,"r",encoding="utf-8") as f:
        t_string = f.read()
    with open(Config.color_table_path, "r", encoding="utf-8") as f:
        t_color = f.read()
    # print(pred_str,pred_color)
    strings = [search_table(t_string,item)for item in pred_str[0]]
    colors = [search_table(t_color,item)for item in pred_color[0]]
    #print (strings,colors)
    return strings,colors


def search_table(table,idx):
    '''
    排除掉小于0或者大于分类数目的索引
    :param table: srting or color
    :param idx: 索引
    :return:
    '''
    if idx < 0:
        return "0"
    elif idx >= len(table):
        return "1"
    else:
        return table[idx]


if __name__ == "__main__":
    #thread_01 = threading.Thread(target=predict, args=["1.png"])
    #thread_02 = threading.Thread(target=predict, args=["2.png"])
    #thread_01.start()
    #thread_02.start()
    #root_dir = r"D:\tensorflow\发票crnn\验证码\img"
    #for m in os.listdir(root_dir):
    #    n = os.path.join(root_dir, m)
    #
    #    strings,colors=predict(n)
    strings,colors=predict("3.png")
    #print (strings,colors)